#include "structMeshFunctions.hpp"
#include "macros.hpp"
#include <iostream>

namespace hms
{

template<typename MeshType, typename Callable, typename ... Ts>
void traverseEdges(
	[[maybe_unused]] EdgeTraversal traversal,
	const MeshBase<MeshType>& mesh,
	Callable func,
	Ts&& ... args
){
	static_assert( std::is_invocable_r_v
		<void, Callable, const MeshType&, Index, Ts ...>,
		"Callable must be invocable with const MeshType& as first "
		"and Index as second argument!"
	);

	#define FORWARD_ARGS mesh, func, std::forward<Ts>(args) ...

	using T = EdgeTraversal;

	if constexpr ( std::is_base_of_v<StructMeshBase<MeshType>, MeshType> ){
		switch (traversal){
			case T::ordered:
				orderedTraversal( FORWARD_ARGS );
				break;

			case T::rowwise:
				rowwiseTraversal( FORWARD_ARGS );
				break;

			case T::blockwise:
				blockwiseTraversal( FORWARD_ARGS );
				break;

			case T::blockwise_alt:
				blockwiseTraversalAlt( FORWARD_ARGS );
				break;

			default:
				std::cerr << "WARNING: Unknown edge traversal type. "
					"Defaulting to ordered traversal.\n";
				orderedTraversal( FORWARD_ARGS );
				break;
		}
	} else {
		assert( traversal == T::ordered &&
			"For unstructured meshes, use ordered edge traversal!" );
		orderedTraversal( FORWARD_ARGS );
	}
}

TRAVERSAL_FUNC_FIRST_ORDER(orderedTraversal){

	OMP_PRAGMA(parallel for schedule(static))
	for (Index i = 0; i<mesh.nEdges(); ++i)
		func( mesh, i, std::forward<Ts>(args) ... );
}

TRAVERSAL_FUNC_FIRST_ORDER(rowwiseTraversal){
	STRUCTURED_TRAVERSAL_ASSERT;
	auto call = [&](Index i){
		func( mesh, i, std::forward<Ts>(args) ... );
	};

	Index
		nx { mesh.derived().size().x() },
		ny { mesh.derived().size().y() };

	Index nh { nHorizontalEdges( mesh.derived().size() ) };

	OMP_PRAGMA(parallel for schedule(static))
	for (int j=0; j<ny; ++j){
		/* bottom and left edges */
		for (int i=0; i<nx; ++i){
			call( i + j*nx );
			call( i + j*(nx+1) + nh );
		}
		/* right boundary */
		call( nx + j*(nx+1) + nh );
	}
	/* top boundary */
	for (int i=0; i<nx; ++i)
		call( i + ny*nx );
}

TRAVERSAL_FUNC_FIRST_ORDER(blockwiseTraversal){
	STRUCTURED_TRAVERSAL_ASSERT;

	auto call = [&](Index i){
		func( mesh, i, std::forward<Ts>(args) ... );
	};

	/* square root of number of cells per block */
	Index blocksize { mesh.derived().blocksize() };

	Index
		nx { mesh.derived().size().x() },
		nh { nHorizontalEdges( mesh.derived().size() ) };

	Array2i nBlocks, rem;
	for (int i=0; i<2; ++i){
		nBlocks[i] = mesh.derived().size()[i] / blocksize;
		rem    [i] = mesh.derived().size()[i] % blocksize;
	}

	/** Each block looks like this:
	 * e.g. blocksize = 3:
	 * 
	 * |_|_|_
	 * |_|_|_
	 * |_|_|_
	 * 
	 * so, for it to be tileable, the right and top edges are not included.
	 * */
	OMP_PRAGMA(parallel)
	{
	OMP_PRAGMA( for schedule(static) collapse(2) nowait )
	for (int j=0; j<nBlocks.y(); ++j){
	for (int i=0; i<nBlocks.x(); ++i){
		Index
			firstH { i*blocksize + j*blocksize*nx },
			firstV { i*blocksize + j*blocksize*(nx+1) + nh };
		for (Index v=0; v<blocksize; ++v){
		for (Index h=0; h<blocksize; ++h){
			call( firstH + v* nx    + h );
			call( firstV + v*(nx+1) + h );
		}}
	}}

	/* remainders */
	/* right side, including the right boundary, e.g. for blocksize = 3 and
	 * 2 = remainder in x:
	 * 
	 * |_|_|
	 * |_|_|
	 * |_|_|
	 * 
	 * */
	OMP_PRAGMA(for schedule(static) nowait)
	for (int j=0; j<nBlocks.y(); ++j){
		Index
			firstH { nBlocks.x()*blocksize + j*blocksize*nx },
			firstV { nBlocks.x()*blocksize + j*blocksize*(nx+1) + nh };
		for (Index v=0; v<blocksize; ++v){
			/* remaining block, same height as before, but smaller in width */
			for (Index h=0; h<rem.x(); ++h){
				call( firstH + v* nx    + h );
				call( firstV + v*(nx+1) + h );
			}
			/* right boundary edge */
			call( firstV + v*(nx+1) + rem.x() );
		}
	}
	/* top, including boundary, e.g. for blocksize = 3 and 1 being the remainder
	 * in y:
	 *  _ _ _
	 * |_|_|_
	 * 
	 * */
	OMP_PRAGMA(for schedule(static) nowait)
	for (int i=0; i<nBlocks.x(); ++i){
		Index
			firstH { i*blocksize + nBlocks.y()*blocksize*nx },
			firstV { i*blocksize + nBlocks.y()*blocksize*(nx+1) + nh };
		for (Index v=0; v<rem.y(); ++v){
			/* remaining block, same width as before, but smaller in height */
			for (Index h=0; h<blocksize; ++h){
				call( firstH + v* nx    + h );
				call( firstV + v*(nx+1) + h );
			}
		}
		/* top boundary edges */
		for (Index h=0; h<blocksize; ++h)
			call( firstH + rem.y()*nx + h );
	}

	/* top right corner, e.g. for remainder of 2 in x and 1 in y:
	 *	_ _
	 * |_|_|
	 * 
	 *  */
	OMP_PRAGMA( single )
	{
	Index
		firstH { nBlocks.x()*blocksize + nBlocks.y()*blocksize*nx },
		firstV { nBlocks.x()*blocksize + nBlocks.y()*blocksize*(nx+1) + nh };
	for (Index v=0; v<rem.y(); ++v){
		/* block remainder */
		for (Index h=0; h<rem.x(); ++h){
			call( firstH + v* nx    + h );
			call( firstV + v*(nx+1) + h );
		}
		/* right boundary edge */
		call( firstV + v*(nx+1) + rem.x() );
	}
	/* top boundary edges */
	for (Index h=0; h<rem.x(); ++h)
		call( firstH + rem.y()*nx + h );

	}}

}

TRAVERSAL_FUNC_FIRST_ORDER(blockwiseTraversalAlt){
	STRUCTURED_TRAVERSAL_ASSERT;

	auto call = [&](Index i){
		func( mesh, i, std::forward<Ts>(args) ... );
	};

	/* square root of number of cells per block */
	Index blocksize { mesh.derived().blocksize() };

	Index
		nx { mesh.derived().size().x() },
		ny { mesh.derived().size().y() },
		nh { nHorizontalEdges( mesh.derived().size() ) };

	Array2i nBlocks, rem;
	for (int i=0; i<2; ++i){
		nBlocks[i] = mesh.derived().size()[i] / blocksize;
		rem    [i] = mesh.derived().size()[i] % blocksize;
	}

	/* remainder blocks, right side */
	for(Index by = 0; by < nBlocks[1]; by++){
		for(Index col = nx-rem[0]; col<nx; col++){
			call(nh + by*blocksize*(nx+1) + col+1);
		}
		for(Index row = by*blocksize+1; row<(by+1)*blocksize; row++){
			for(Index col = nx-rem[0]; col<nx; col++){
				call(row*nx + col);
				call(nh + row*(nx+1)+col+1);
			}
		}
	}

	/* remainder blocks, top */
	for(Index bx = 0; bx < nBlocks[0]; bx++){
		for(Index row = ny-rem[1]; row<ny; row++){
			call((row+1)*nx + bx*blocksize);
			for(Index col = bx*blocksize+1; col<(bx+1)*blocksize; col++){
				call((row+1)*nx + col);
				call(row*(nx+1) + col);
			}
		}
	}

	/* remainder block, top right corner */
	for(Index row = ny-rem[1]; row<ny; row++){
		for(Index col=nx-rem[0]; col< nx; col++){
			call((row+1)*nx + col);
			call(nh + row*(nx+1) + col + 1);
		}
	}

	/* blocks */
	for(Index bx = 0; bx<nBlocks[0]; bx++){
		for(Index by = 0; by<nBlocks[1]; by++){
			for(Index col=bx*blocksize+1; col<(bx+1)*blocksize; col++){
				call(nh+(by*blocksize)*(nx+1)+col);
			}
			for(Index row=by*blocksize+1; row<(by+1)*blocksize; row++){
				call(row*nx + bx*blocksize);
				for(Index col = bx*blocksize+1; col<(bx+1)*blocksize; col++){
					call(row*nx+col);
					call(nh+row*(nx+1)+col);
				}
			}
		}
	}

	/* shared edges */
	// this is threadsafe for blocksize >= 1, but will probably produce more cache misses
	for(Index bx = 0; bx<=nBlocks[0]; bx++){
		for(Index row = 0; row<ny; row++){
			call(nh + row*(nx+1) + bx*blocksize);
		}
	}

	for(Index by = 0; by<=nBlocks[1]; by++){
		for(Index col = 0; col<nx; col++){
			call(by*blocksize*nx + col);
		}
	}
}

#undef FORWARD_ARGS

} // end namespace hms